import torch
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import seaborn as sns
import cv2


class ComprehensiveVisualizer:
    """综合可视化器，整合多种可视化功能"""
    
    def __init__(self, save_dir: Path):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        self.hooks = []
        self.activations = {}
        
    def register_hooks(self, model: torch.nn.Module, layer_names: List[str]):
        """注册钩子捕获激活"""
        def get_activation(name):
            def hook(model, input, output):
                if isinstance(output, torch.Tensor):
                    self.activations[name] = output.detach().cpu()
            return hook
        
        for name, module in model.named_modules():
            if name in layer_names:
                hook = module.register_forward_hook(get_activation(name))
                self.hooks.append(hook)
    
    def remove_hooks(self):
        """移除所有钩子"""
        for hook in self.hooks:
            hook.remove()
        self.hooks.clear()
    
    def visualize_model_summary(self, model: torch.nn.Module):
        """可视化模型概要信息"""
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        
        # 统计不同类型的层
        layer_types = {}
        for name, module in model.named_modules():
            module_type = type(module).__name__
            if module_type not in layer_types:
                layer_types[module_type] = 0
            layer_types[module_type] += 1
        
        # 创建可视化
        fig, axes = plt.subplots(1, 2, figsize=(15, 6))
        fig.suptitle('Model Summary', fontsize=16)
        
        # 参数统计
        param_data = ['Total Parameters', 'Trainable Parameters', 'Non-trainable Parameters']
        param_values = [total_params, trainable_params, total_params - trainable_params]
        
        axes[0].bar(param_data, param_values)
        axes[0].set_title('Parameter Statistics')
        axes[0].set_ylabel('Number of Parameters')
        for i, v in enumerate(param_values):
            axes[0].text(i, v, f'{v:,}', ha='center', va='bottom')
        
        # 层类型分布
        layer_names = list(layer_types.keys())[:10]  # 显示前10种类型
        layer_counts = [layer_types[name] for name in layer_names]
        
        axes[1].pie(layer_counts, labels=layer_names, autopct='%1.1f%%')
        axes[1].set_title('Layer Type Distribution')
        
        plt.tight_layout()
        save_path = self.save_dir / 'model_summary.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        return {
            'total_parameters': total_params,
            'trainable_parameters': trainable_params,
            'model_size_mb': total_params * 4 / 1024 / 1024,  # 假设float32
            'layer_types': layer_types
        }
    
    def visualize_activation_patterns(self, layer_name: str, max_channels: int = 16):
        """可视化激活模式"""
        if layer_name not in self.activations:
            print(f"Layer {layer_name} not found in activations")
            return
        
        activation = self.activations[layer_name]
        if len(activation.shape) != 4:  # 需要 (B, C, H, W)
            print(f"Activation shape {activation.shape} not supported for visualization")
            return
        
        batch_size, channels, height, width = activation.shape
        num_channels = min(channels, max_channels)
        
        cols = 4
        rows = (num_channels + cols - 1) // cols
        fig, axes = plt.subplots(rows, cols, figsize=(16, 4 * rows))
        fig.suptitle(f'Activation Patterns - {layer_name}', fontsize=16)
        
        if rows == 1:
            axes = axes.reshape(1, -1)
        
        for i in range(num_channels):
            row = i // cols
            col = i % cols
            
            channel_activation = activation[0, i].numpy()
            
            im = axes[row, col].imshow(channel_activation, cmap='viridis')
            axes[row, col].set_title(f'Channel {i}')
            axes[row, col].axis('off')
            plt.colorbar(im, ax=axes[row, col], fraction=0.046, pad=0.04)
        
        # 隐藏空的子图
        for i in range(num_channels, rows * cols):
            row = i // cols
            col = i % cols
            axes[row, col].axis('off')
        
        plt.tight_layout()
        save_path = self.save_dir / f'activation_patterns_{layer_name.replace(".", "_")}.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def visualize_layer_statistics(self):
        """可视化层统计信息"""
        if not self.activations:
            print("No activations available for statistics")
            return
        
        layer_stats = {}
        for layer_name, activation in self.activations.items():
            flat_activation = activation.flatten().numpy()
            layer_stats[layer_name] = {
                'mean': np.mean(flat_activation),
                'std': np.std(flat_activation),
                'sparsity': np.mean(flat_activation == 0),
                'max': np.max(flat_activation),
                'min': np.min(flat_activation)
            }
        
        # 创建统计图表
        layer_names = list(layer_stats.keys())
        stats_names = ['mean', 'std', 'sparsity']
        
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        fig.suptitle('Layer Statistics', fontsize=16)
        
        for idx, stat_name in enumerate(stats_names):
            row = idx // 2
            col = idx % 2
            
            values = [layer_stats[name][stat_name] for name in layer_names]
            
            axes[row, col].bar(range(len(layer_names)), values)
            axes[row, col].set_title(f'{stat_name.title()} by Layer')
            axes[row, col].set_xticks(range(len(layer_names)))
            axes[row, col].set_xticklabels([name.split('.')[-1][:8] for name in layer_names], 
                                         rotation=45)
        
        # 激活范围
        ranges = [layer_stats[name]['max'] - layer_stats[name]['min'] for name in layer_names]
        axes[1, 1].bar(range(len(layer_names)), ranges)
        axes[1, 1].set_title('Activation Range by Layer')
        axes[1, 1].set_xticks(range(len(layer_names)))
        axes[1, 1].set_xticklabels([name.split('.')[-1][:8] for name in layer_names], 
                                 rotation=45)
        
        plt.tight_layout()
        save_path = self.save_dir / 'layer_statistics.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        return layer_stats
    
    def create_comprehensive_report(self, model_info: Dict, layer_stats: Dict):
        """创建综合分析报告"""
        report_path = self.save_dir / "visualization_report.txt"
        
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write("=== Model Visualization Report ===\n\n")
            
            f.write("1. Model Overview:\n")
            f.write(f"   - Total Parameters: {model_info['total_parameters']:,}\n")
            f.write(f"   - Trainable Parameters: {model_info['trainable_parameters']:,}\n")
            f.write(f"   - Model Size: {model_info['model_size_mb']:.2f} MB\n\n")
            
            f.write("2. Layer Analysis:\n")
            if layer_stats:
                f.write(f"   - Number of analyzed layers: {len(layer_stats)}\n")
                avg_sparsity = np.mean([stats['sparsity'] for stats in layer_stats.values()])
                f.write(f"   - Average sparsity: {avg_sparsity:.4f}\n")
            
            f.write("\n3. Generated Visualizations:\n")
            f.write("   - model_summary.png: Model parameter and layer distribution\n")
            f.write("   - activation_patterns_*.png: Layer activation patterns\n")
            f.write("   - layer_statistics.png: Statistical analysis of layers\n")
        
        print(f"Visualization report saved to: {report_path}")


class ModelInterpreter:
    """模型解释器，提供模型架构分析"""
    
    def __init__(self, model: torch.nn.Module, save_dir: Path):
        self.model = model
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
    
    def visualize_model_architecture(self):
        """可视化模型架构"""
        visualizer = ComprehensiveVisualizer(self.save_dir)
        return visualizer.visualize_model_summary(self.model)


class LayerAnalyzer:
    """层分析器"""
    
    def __init__(self, save_dir: Path):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        self.hooks = []
        self.layer_outputs = {}
    
    def register_hooks(self, model: torch.nn.Module, layer_names: List[str]):
        """注册钩子"""
        def get_output(name):
            def hook(model, input, output):
                if isinstance(output, torch.Tensor):
                    self.layer_outputs[name] = output.detach().cpu()
            return hook
        
        for name, module in model.named_modules():
            if name in layer_names:
                hook = module.register_forward_hook(get_output(name))
                self.hooks.append(hook)
    
    def remove_hooks(self):
        """移除钩子"""
        for hook in self.hooks:
            hook.remove()
        self.hooks.clear()
    
    def analyze_layer_statistics(self):
        """分析层统计信息"""
        visualizer = ComprehensiveVisualizer(self.save_dir)
        visualizer.activations = self.layer_outputs
        return visualizer.visualize_layer_statistics()
